open Core.Std
open Printf
open Semantics (* for port *)

let net_name i1 i2 i3 =
  sprintf "n%d_%d_%d" i1 i2 i3

let reduce_layer (n : int) (layer_index : int) (layer : int array) =
  let compute_triple m =
    let q = m / 3 in
    let r = m mod 3 in
    match r with
    | 0 -> 0, 0, q
    | 1 -> 1, 0, q
    | 2 -> 0, 1, q
    | _ -> assert false
  in
  let src_net_name = net_name layer_index in
  let dst_net_name = net_name (layer_index+1) in
  let next_layer = Array.create 0 ~len:n in
  let layer_assigns = ref [] in
  let layer_instances = ref [] in
  let triple_arr = Array.map layer ~f:compute_triple in
  for i = n-1 downto 0 do
    let n1, n2, n3 = triple_arr.(i) in
    let src_bit = ref 0 in
    let dst_bit = ref 0 in
    let assigns =
      List.map (List.range 0 n1) ~f:begin fun _ ->
        let src_name = sprintf "n%d_%d_%d" layer_index i !src_bit in
        let dst_name = sprintf "n%d_%d_%d" (layer_index+1) i !dst_bit in
        src_bit := !src_bit+1;
        dst_bit := !dst_bit+1;
        (dst_name, (0, 1)), Vlog.Sym src_name
      end
    in
    let adder_name () = sprintf "a%d_%d_%d" (layer_index+1) i !dst_bit in
    let gen_adders is_full num =
      List.map (List.range 0 num) ~f:begin fun _ ->
        let inst_name = adder_name () in
        let (ports : (string * port * Vlog.expr) list) =
          let input_ports_2 =
            [
              "b", (1, Input ), Vlog.Sym (src_net_name i (!src_bit+1));
              "a", (1, Input ), Vlog.Sym (src_net_name i (!src_bit  ));
            ]
          in
          let input_ports_all =
            if is_full
            then begin
              let cin_port =
                "cin", (1, Input ), Vlog.Sym (src_net_name i (!src_bit+2))
              in
              cin_port :: input_ports_2
            end
            else input_ports_2
          in
          let out_port =
            "out", (1, Output), Vlog.Sym (dst_net_name i !dst_bit)
          in
          src_bit := !src_bit + (if is_full then 3 else 2);
          dst_bit := !dst_bit + 1;
          if i+1 < n
          then begin
            let cout_port =
              "cout", (1, Output), Vlog.Sym (dst_net_name (i+1) next_layer.(i+1))
            in
            next_layer.(i+1) <- next_layer.(i+1) + 1;
            cout_port :: out_port :: input_ports_all
          end
          else out_port :: input_ports_all
        in
        inst_name,
        { Vlog.class_name = if is_full then "full_adder" else "half_adder";
          parameters = [];
          ports }
      end
    in
    let half_adders = gen_adders false n2 in
    let full_adders = gen_adders true n3 in
    next_layer.(i) <- !dst_bit;
    layer_assigns := assigns @ !layer_assigns;
    layer_instances := half_adders @ full_adders @ !layer_instances;
  done;
  next_layer, !layer_assigns, !layer_instances

let main () =
  let n = 32 in
  let layer0 = Array.create 0 ~len:n in
  for i = 0 to n-1 do
    layer0.(i) <- i+1
  done;
  let initial_assigns =
    List.map (List.range 0 n) ~f:begin fun weight ->
      List.map (List.range 0 (weight+1)) ~f:begin fun i ->
        (net_name 0 weight i, (0, 1)),
        Vlog.Binary begin
          Vlog.And,
          Vlog.Sym_part ("a", (i, i+1)),
          Vlog.Sym_part ("b", (weight-i, weight-i+1))
        end
      end
    end
    |> List.concat
  in
  let layer_index = ref 0 in
  let layers = ref [layer0, initial_assigns, []] in
  let last_layer () =
    let layer, _, _ = List.hd_exn !layers in
    layer
  in
  while not (Array.for_all (last_layer ()) ~f:(fun m -> m <= 2)) do
    let next_layer, layer_assigns, layer_instances =
      reduce_layer n !layer_index (last_layer ())
    in
    layers := (next_layer, layer_assigns, layer_instances) :: !layers;
    layer_index := !layer_index + 1;
  done;
  let final_layer = last_layer () in
  let final_layer_index = !layer_index in
  let k =
    let i = ref 0 in
    while final_layer.(!i) = 1 do
      i := !i+1
    done;
    !i
  in
  let output_assign_direct =
    ("out", (0, k)),
    Vlog.Concat begin
      List.map (List.range 0 k)
      ~f:(fun i -> Vlog.Sym (net_name final_layer_index i 0))
    end
  in
  let output_assign_cpa =
    ("out", (k, n)),
    Vlog.Binary begin
      Vlog.Add,
      Vlog.Concat begin
        List.range k n
        |> List.rev
        |> List.map ~f:(fun i -> Vlog.Sym (net_name final_layer_index i 0))
      end,
      Vlog.Concat begin
        List.range k n
        |> List.rev
        |> List.map ~f:(fun i -> Vlog.Sym (net_name final_layer_index i 1))
      end
    end
  in
  let layers = List.rev !layers in
  let wire_alist_of_layer layer_index layer =
    List.mapi (Array.to_list layer) ~f:begin fun i m ->
      List.map (List.range 0 m) ~f:(fun j -> net_name layer_index i j, 1)
    end
    |> List.concat
  in
  let wire_alist =
    List.mapi layers ~f:(fun i (layer, _, _) -> wire_alist_of_layer i layer)
    |> List.concat
  in
  let layer_assigns =
    List.map layers ~f:(fun (_, a, _) -> a)
    |> List.concat
  in
  let instance_alist =
    List.map layers ~f:(fun (_, _, a) -> a)
    |> List.concat
  in
  let port_alist =
    [ "a", (32, Input);
      "b", (32, Input);
      "out", (32, Output) ]
  in
  let vlog_module =
    { Vlog.port_alist;
      wire_alist;
      assigns = output_assign_cpa :: output_assign_direct :: layer_assigns;
      instance_alist;
      reg_alist = [];
      has_clock = false }
  in
  Vlog.write_named_vlog_module ("mul", vlog_module) "mul.v"

let () =
  main ()
